import cv2
import numpy as np
class DetRectToOCRPolygon():
    """
    图片resize, 可以用于避免输入图像过大导致爆显存等问题

    args:
        keys: 指定了从data中读取数据的关键字和保存数据的关键字
            格式: keys={"in":"your_in_keyword",
                        "out":"your_save_keyword"
                        }
    """
    def __init__(self, keys, target_cls):
        self.keys = keys
        self.target_cls = set(target_cls)

    def __call__(self, data:dict):
        det_result = data[self.keys["in"]]
        ocr_polygons = []
        for boxes in det_result["boxes"]:
            polygons = []
            for box in boxes:
                cls_id, score, x1,y1,x2,y2 = box
                if cls_id in self.target_cls:
                    box_ = [x1,y1,x2,y1,x2,y2,x1,y2,score]
                    polygons.append(box_)
            ocr_polygons.append(polygons)
        data[self.keys["out"]] = ocr_polygons
        return data